'''
Author: WuFeng <763467339@qq.com>
Date: 2024-11-20 14:29:04
LastEditTime: 2024-11-20 19:08:36
LastEditors: WuFeng <763467339@qq.com>
Description: pdf密码相关
FilePath: \markdown-magic\python\service\pdfPassword.py
Copyright 版权声明
'''
import PyPDF2

def is_pdf_encrypted(pdf_path):
    """
    检查 PDF 文件是否加密。
    :param pdf_path: PDF 文件的路径。
    :return: 如果 PDF 文件加密则返回 True，否则返回 False。
    """
    with open(pdf_path, 'rb') as file:
        reader = PyPDF2.PdfReader(file)
        return reader.is_encrypted


def remove_pdf_password(input_pdf, output_pdf, password):
    """
    解除 PDF 文件的密码保护。

    :param input_pdf: 输入的带密码的 PDF 文件路径。
    :param output_pdf: 输出的无密码的 PDF 文件路径。
    :param password: PDF 文件的密码。
    """
    # 打开输入 PDF 文件
    with open(input_pdf, 'rb') as infile:
        reader = PyPDF2.PdfReader(infile)
        
        # 检查 PDF 是否加密
        if not reader.is_encrypted:
            print("输入的 PDF 文件未加密。")
            return 'noPassword'
        
        # 尝试解密 PDF 文件
        if not reader.decrypt(password):
            print("密码错误，无法解密 PDF 文件。")
            return 'passwordError'
        
        # 创建输出 PDF 文件
        writer = PyPDF2.PdfWriter()
        
        # 将所有页面添加到输出 PDF 文件
        for page in reader.pages:
            writer.add_page(page)
        
        # 写入输出 PDF 文件
        with open(output_pdf, 'wb') as outfile:
            writer.write(outfile)
    
    print(f"PDF 文件已成功解密并保存为 {output_pdf}")
    return 'success'


def set_pdf_password(input_pdf, output_pdf, user_password, owner_password):
    """
    为 PDF 文件设置密码保护。

    :param input_pdf: 输入的无密码的 PDF 文件路径。
    :param output_pdf: 输出的带密码的 PDF 文件路径。
    :param user_password: 用户密码，用于打开 PDF 文件。
    :param owner_password: 所有者密码，用于修改 PDF 文件的权限。
    """
    # 打开输入 PDF 文件
    with open(input_pdf, 'rb') as infile:
        reader = PyPDF2.PdfReader(infile)
        
        # 创建输出 PDF 文件
        writer = PyPDF2.PdfWriter()
        
        # 将所有页面添加到输出 PDF 文件
        for page in reader.pages:
            writer.add_page(page)
        
        # 设置密码
        writer.encrypt(user_password, owner_password, use_128bit=True)
        
        # 写入输出 PDF 文件
        with open(output_pdf, 'wb') as outfile:
            writer.write(outfile)
    
    print(f"PDF 文件已成功设置密码并保存为 {output_pdf}")
    return 'success'

# # 使用示例
# input_pdf = 'path/to/your/unencrypted.pdf'
# output_pdf = 'path/to/your/encrypted.pdf'
# user_password = 'user_password'
# owner_password = 'owner_password'

# set_pdf_password(input_pdf, output_pdf, user_password, owner_password)